#include "det_predictor.h"

namespace ai {


ONNXDetPredictor::ONNXDetPredictor(ONNXPackPredictor* model, ONNXModelManager* manager, LogInfo *lg):
  ONNXPredictor(model, manager, lg) { }


ONNXDetPredictor::~ONNXDetPredictor(){

}


int ONNXDetPredictor::RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) {
  int ret = 0;
  
  stream::ImageBlob img_blob(stream::ImageBlobMode_BGR);
  if (!mdl_ox->cfg->transforms->run(blob->img, img_blob, mdl_ox->cfg)) { 
    printf("transforms->run fail \n");
    return model_image_channels_check_error;
  }
  // std::cout << img_blob.img.cols << ", " << img_blob.img.rows << ", c:" << img_blob.img.channels() << std::endl;
  // std::cout << "img_blob.scale: " << img_blob.scale << std::endl;
  // img_blob.img.convertTo(img_blob.img, CV_32FC(img_blob.img.channels()));
  // std::cout << img_blob.img.type() << std::endl;
  
  std::vector<Ort::Value> inputTensors;
  std::vector<const char *> input_names;
  for (int idx = 0; idx < mdl_ox->cfg->input_names.size(); idx++) {
    auto& input_name = mdl_ox->cfg->input_names[idx];
    input_names.push_back(input_name.c_str());
    
    if (input_name == "image" || input_name == "images" || input_name == "input") {
      std::vector<int64_t> shape;
      float* pdata = nullptr;
      if (mdl_ox->cfg->data_format == "CHW"){
        pdata = img_blob.im_vec_data.data();
        shape = { 1, mdl_ox->cfg->channels, img_blob.new_im_shape[0], img_blob.new_im_shape[1] }; 
      }
      else { 
        pdata = (float*)img_blob.img.data;
        shape = { 1, img_blob.new_im_shape[0], img_blob.new_im_shape[1], mdl_ox->cfg->channels }; 
      }
      // std::cout << "shape 0: " << shape[0] << ", "<< shape[1] << ", " << shape[2]<< ", " << shape[3] << std::endl;

      // float* pdata_ = (float*)img_blob.img.data;
      // std::cout << "pdata 0:" << pdata[0] << ", " << pdata[1] << ", " << pdata[2]<< ", " << pdata[3]<< ", " << pdata[4]<< ", " << pdata[5]<< ", " << pdata[6]<< ", " << pdata[7]<< ", " << pdata[8]<< ", " << pdata[9] << ", " << pdata[10] << ", " << pdata[11]<< ", " << pdata[12]<< ", " << pdata[13]<< ", " << pdata[14]<< ", " << pdata[15]<< ", " << pdata[16]<< ", " << pdata[17] << std::endl;
      // std::cout << "pdata 200:" << pdata[128000+0] << ", " << pdata[128000+1] << ", " << pdata[128000+2]<< ", " << pdata[128000+3]<< ", " << pdata[128000+4]<< ", " << pdata[128000+5]<< ", " << pdata[128000+6]<< ", " << pdata[128000+7]<< ", " << pdata[128000+8]<< ", " << pdata[128000+9] << ", " << pdata[128000+10] << ", " << pdata[128000+11]<< ", " << pdata[128000+12]<< ", " << pdata[128000+13]<< ", " << pdata[128000+14]<< ", " << pdata[128000+15]<< ", " << pdata[128000+16]<< ", " << pdata[128000+17] << std::endl;
      // std::cout << "pdata 400:" << pdata[256000+0] << ", " << pdata[256000+1] << ", " << pdata[256000+2]<< ", " << pdata[256000+3]<< ", " << pdata[256000+4]<< ", " << pdata[256000+5]<< ", " << pdata[256000+6]<< ", " << pdata[256000+7]<< ", " << pdata[256000+8]<< ", " << pdata[256000+9] << ", " << pdata[256000+10] << ", " << pdata[256000+11]<< ", " << pdata[256000+12]<< ", " << pdata[256000+13]<< ", " << pdata[256000+14]<< ", " << pdata[256000+15]<< ", " << pdata[256000+16]<< ", " << pdata[256000+17] << std::endl;
      inputTensors.push_back(Ort::Value::CreateTensor<float>(mdl_ox->ortMemoryInfo, 
                        pdata, 
                        accumulate(shape.begin(), shape.end(), 1, std::multiplies<int>()), 
                        shape.data(), 
                        shape.size()));
    }
    else if (input_name == "scale_factor") {
      std::vector<int64_t> shape = {1,2};
      std::vector<float> os = {img_blob.scale, img_blob.scale};
      // std::vector<float> os = {0.00001, 0.00001};
      inputTensors.push_back(Ort::Value::CreateTensor<float>(mdl_ox->ortMemoryInfo, 
                        os.data(), 
                        accumulate(shape.begin(), shape.end(), 1, std::multiplies<int>()), 
                        shape.data(), 
                        shape.size()));
    }

  }


  std::vector<const char *> output_names;
  for (auto& str : mdl_ox->cfg->output_names) {output_names.push_back(str.c_str());}

  // double t1 = cv::getTickCount();;
  std::vector<Ort::Value> output_tensor_values = mdl_ox->session->Run(
        Ort::RunOptions{nullptr}, input_names.data(),
        inputTensors.data(), input_names.size(), 
        output_names.data(), output_names.size());

  // double tt = ((double)cv::getTickCount() - t1) / cv::getTickFrequency();
  // std::cout << " time:" << tt/100 << std::endl;
  // std::cout << " output_tensor_values.size():" << output_tensor_values.size() << std::endl;

  outputs.clear();
  outputs_size.clear();
  for (size_t idx = 0; idx < output_tensor_values.size(); idx++) {
    // size_t count = output_tensor_values[idx].IsSparseTensor();
    // std::cout << " IsSparseTensor:" << output_tensor_values[idx].IsSparseTensor() << std::endl;
    // std::cout << " IsTensor:" << output_tensor_values[idx].IsTensor() << std::endl;
    int output_size = 0;
    if (output_tensor_values[idx].IsTensor()) {
      Ort::TensorTypeAndShapeInfo info = output_tensor_values[idx].GetTensorTypeAndShapeInfo();
      
      std::vector<int64_t> shape = info.GetShape();
      output_size = accumulate(shape.begin(), shape.end(), 1, std::multiplies<int>());
      // std::cout << "output_tensor_values shape: ";
      // for (int ii = 0; ii < shape.size(); ii++) { std::cout << " " << shape[ii]; }
      // std::cout << std::endl;
    }


    float* cur_data = output_tensor_values[idx].GetTensorMutableData<float>();
    // std::cout << cur_data[0] << "  " << cur_data[1] << "  " << (double)cur_data[2] << "  " << (double)cur_data[3]<< "  " << (double)cur_data[4] << "  " << (double)cur_data[5] << "  " << cur_data[6] << "  " << cur_data[7] << "  " << (double)cur_data[8] << "  " << (double)cur_data[9] << "  " << (double)cur_data[10] << "  " << (double)cur_data[11]<< std::endl;

    outputs.push_back(cur_data);
    outputs_size.push_back(output_size);
  }
  
  if (mdl_ox->cfg->arch == "YOLOV5.paddle") {
    GetConvDetectionResult_paddle(img_blob, det_infos);
  }
  else if (mdl_ox->cfg->arch == "YOLOV5.plate") {
    GetConvDetectionResult_plate(img_blob, det_infos);
  }
  else if (mdl_ox->cfg->arch == "YOLOV8.out1") {
    GetConvDetectionResult(img_blob, det_infos);
  }
  else if (mdl_ox->cfg->arch == "YOLOV8.nosigmoid.out1") {
    GetConvDetectionResult_nosigmoid(img_blob, det_infos);
  } else {
    spdlog::get("logger")->info("ERROR. ONNXDetPredictor::RunDet. mdl_ox->cfg->arch error: {} ", mdl_ox->cfg->arch);
  }

  outputs.clear();
  outputs_size.clear();

  // std::cout << " NMSBoxes:" << det_infos.size() << std::endl;
  return 0;
}

int ONNXDetPredictor::GetConvDetectionResult(stream::ImageBlob& img_blob, std::vector<BaseInfo*>& det_infos) {
  int ret = 0;

  std::vector<cv::Rect> boxs;
  std::vector<float> confidences;
  std::vector<int> labels;


  int class_num = mdl_ox->cfg->label_list.size();
  int net_width = class_num+4;
  
  for (int idx = 0; idx < outputs.size(); idx++) {
	  cv::Mat output0 = cv::Mat(cv::Size((int)mdl_ox->output_tensor_dims[idx][2], (int)mdl_ox->output_tensor_dims[idx][1]), CV_32F, outputs[idx]).t();  //[bs,116,8400]=>[bs,8400,116]
    float* pdata = (float*)output0.data;
    int rows = output0.rows;

    for (int r = 0; r < rows; ++r) {    //stride
      cv::Mat scores(1, class_num, CV_32F, pdata + 4);
      cv::Point classIdPoint;
      double max_class_socre;
      cv::minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
      max_class_socre = (float)max_class_socre;
      if (max_class_socre >= mdl_ox->cfg->draw_threshold) {

        //rect [x,y,w,h]
        float x = pdata[0] / img_blob.scale;
        float y = pdata[1] / img_blob.scale;
        int w = pdata[2] / img_blob.scale;
        int h = pdata[3] / img_blob.scale;
        // std::cout << "score:" << max_class_socre << " idx:" << classIdPoint.x << " x:" << x << " y: " << y << " w: " << w << " h:" << h << std::endl;

        int left = std::max(int(x - 0.5 * w + 0.5), 0);
        int top = std::max(int(y - 0.5 * h + 0.5), 0);
        if (left + w >= img_blob.ori_im_shape[1]) { w = img_blob.ori_im_shape[1] - left - 1;}
        if (top + h >= img_blob.ori_im_shape[0]) { h = img_blob.ori_im_shape[0] - top - 1;}
        labels.push_back(classIdPoint.x);
        confidences.push_back(max_class_socre);
        boxs.push_back(cv::Rect(left, top, w, h));
      }
      pdata += net_width;
    }
  }

  std::vector<int> indices;
  cv::dnn::NMSBoxes(boxs, confidences, mdl_ox->cfg->draw_threshold, mdl_ox->cfg->nms_threshold, indices, 100);
  for (auto& idx : indices) {
    DetInfo* di = new DetInfo(confidences[idx], labels[idx], boxs[idx], mdl_ox->cfg->label_list[labels[idx]]);
    det_infos.emplace_back(di);
  }

  return ret;
}
int ONNXDetPredictor::GetConvDetectionResult_nosigmoid(stream::ImageBlob& img_blob, std::vector<BaseInfo*>& det_infos) {
  int ret = 0;

  std::vector<cv::Rect> boxs;
  std::vector<float> confidences;
  std::vector<int> labels;


  int class_num = mdl_ox->cfg->label_list.size();
  int net_width = class_num+4;
  
  for (int idx = 0; idx < outputs.size(); idx++) {
	  cv::Mat output0 = cv::Mat(cv::Size((int)mdl_ox->output_tensor_dims[idx][2], (int)mdl_ox->output_tensor_dims[idx][1]), CV_32F, outputs[idx]).t();  //[bs,116,8400]=>[bs,8400,116]
    // std::cout << "output0.cols:" << output0.cols << " rows:" << output0.rows << std::endl;

    float* pdata = (float*)output0.data;
    int rows = output0.rows;
    // std::cout << "rows:" << rows << std::endl;

    for (int r = 0; r < rows; ++r) {    //stride
      // std::cout << "pdata:" << pdata[0] << ", " << pdata[1] << ", " << pdata[2] << ", " << pdata[3] << ", " << pdata[4] << ", " << pdata[5] << std::endl;

      cv::Mat scores(1, class_num, CV_32F, pdata + 4);
      cv::Point classIdPoint;
      double max_class_socre;
      cv::minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
      max_class_socre = (float)max_class_socre;
      if (max_class_socre < -1) {
      //   // std::cout << "pdata:" << pdata[0] << ", " << pdata[1] << ", " << pdata[2] << ", " << pdata[3] << ", " << pdata[4] << ", " << pdata[5] << std::endl;
        goto goto_next;
      }
      // std::cout << "pdata:" << pdata[0] << ", " << pdata[1] << ", " << pdata[2] << ", " << pdata[3] << ", " << pdata[4] << ", " << pdata[5] << std::endl;
      max_class_socre = 1.0 / (1+std::exp(-max_class_socre));
      if (max_class_socre >= mdl_ox->cfg->draw_threshold) {
      // std::cout << "pdata:" << pdata[0] << ", " << pdata[1] << ", " << pdata[2] << ", " << pdata[3] << ", " << pdata[4] << ", " << pdata[5] << std::endl;

        //rect [x,y,w,h]
        float x = pdata[0] / img_blob.scale;
        float y = pdata[1] / img_blob.scale;
        int w = pdata[2] / img_blob.scale;
        int h = pdata[3] / img_blob.scale;
        // std::cout << "score:" << max_class_socre << " idx:" << classIdPoint.x << " x:" << x << " y: " << y << " w: " << w << " h:" << h << std::endl;

        int left = std::max(int(x - 0.5 * w + 0.5), 0);
        int top = std::max(int(y - 0.5 * h + 0.5), 0);
        if (left + w >= img_blob.ori_im_shape[1]) { w = img_blob.ori_im_shape[1] - left - 1;}
        if (top + h >= img_blob.ori_im_shape[0]) { h = img_blob.ori_im_shape[0] - top - 1;}
        labels.push_back(classIdPoint.x);
        confidences.push_back(max_class_socre);
        boxs.push_back(cv::Rect(left, top, w, h));
      }
goto_next:
      pdata += net_width;
    }
  }

  std::vector<int> indices;
  cv::dnn::NMSBoxes(boxs, confidences, mdl_ox->cfg->draw_threshold, mdl_ox->cfg->nms_threshold, indices, 100);
  for (auto& idx : indices) {
    DetInfo* di = new DetInfo(confidences[idx], labels[idx], boxs[idx], mdl_ox->cfg->label_list[labels[idx]]);
    det_infos.emplace_back(di);
  }

  return ret;
}

int ONNXDetPredictor::GetConvDetectionResult_plate(stream::ImageBlob& img_blob, std::vector<BaseInfo*>& det_infos) {
  int ret = 0;

  std::vector<cv::Rect> boxs;
  std::vector<std::vector<cv::Point>> land_marks;
  std::vector<float> confidences;
  std::vector<int> labels;
  std::vector<int> double_plates;


  int class_num = mdl_ox->cfg->label_list.size();
  int net_width = 15;
  
  for (int idx = 0; idx < outputs.size(); idx++) {
	  cv::Mat output0 = cv::Mat(cv::Size((int)mdl_ox->output_tensor_dims[idx][2], (int)mdl_ox->output_tensor_dims[idx][1]), CV_32F, outputs[idx]);  
    float* pdata = (float*)output0.data;
    int rows = output0.rows;
    // std::cout << "shape: " << output0.rows << "," << output0.cols << std::endl;

    for (int r = 0; r < rows; ++r) {
      float max_class_socre = 0;
      int double_plate = 0;
      if (pdata[14] > pdata[13]) {
        double_plate = 1;
        max_class_socre = pdata[14] * pdata[4];
      } else {
        max_class_socre = pdata[13] * pdata[4];
      }
      // cv::Mat scores(1, class_num, CV_32F, pdata + 4);
      // cv::Point classIdPoint;
      // double max_class_socre;
      // cv::minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
      // max_class_socre = (float)max_class_socre;
      if (max_class_socre >= mdl_ox->cfg->draw_threshold) {

        //rect [x,y,w,h]
        float w = pdata[2] / img_blob.scale;
        float h = pdata[3] / img_blob.scale;
        float x = pdata[0] / img_blob.scale;
        float y = pdata[1] / img_blob.scale;

        // std::cout << "score:" << max_class_socre << " idx:" << classIdPoint.x << " x:" << x << " y: " << y << " w: " << w << " h:" << h << std::endl;
        // std::cout << pdata[0] << "  " << pdata[1] << "  " << pdata[2] << "  " << pdata[3]<< "  " << pdata[4] << "  " << pdata[5] << "  " << pdata[6] << "  " << pdata[7] << "  " << pdata[8] << "  " << pdata[9] << "  " << pdata[10] << "  " << pdata[11]<< "  " << pdata[12]<< "  " << pdata[13]<< "  " << pdata[14]<< std::endl;
        std::vector<cv::Point> land_mark;
        land_mark.push_back(cv::Point(pdata[5] / img_blob.scale, pdata[6] / img_blob.scale));
        land_mark.push_back(cv::Point(pdata[7] / img_blob.scale, pdata[8] / img_blob.scale));
        land_mark.push_back(cv::Point(pdata[9] / img_blob.scale, pdata[10] / img_blob.scale));
        land_mark.push_back(cv::Point(pdata[11] / img_blob.scale, pdata[12] / img_blob.scale));


        int left = std::max(int(x - 0.5 * w + 0.5), 0);
        int top = std::max(int(y - 0.5 * h + 0.5), 0);
        if (left + w >= img_blob.ori_im_shape[1]) { w = img_blob.ori_im_shape[1] - left - 1;}
        if (top + h >= img_blob.ori_im_shape[0]) { h = img_blob.ori_im_shape[0] - top - 1;}
        double_plates.push_back(double_plate);
        land_marks.push_back(land_mark);
        labels.push_back(0);
        confidences.push_back(max_class_socre);
        boxs.push_back(cv::Rect(left, top, w, h));
      }
      pdata += net_width;
    }
  }

  std::vector<int> indices;
  cv::dnn::NMSBoxes(boxs, confidences, mdl_ox->cfg->draw_threshold, mdl_ox->cfg->nms_threshold, indices, 100);
  for (auto& idx : indices) {
    TextDetInfo* di = new TextDetInfo(land_marks[idx], confidences[idx]);
    di->det_box = boxs[idx];
    di->double_plate = double_plates[idx];
    det_infos.emplace_back(di);
  }

  return ret;
}

int ONNXDetPredictor::GetConvDetectionResult_paddle(stream::ImageBlob& img_blob, std::vector<BaseInfo*>& det_infos) {
  int ret = 0;

  std::vector<cv::Rect> boxs;
  std::vector<float> confidences;
  std::vector<int> labels;


  int class_num = mdl_ox->cfg->label_list.size();
  int step = 6;


  for (int idx = 0; idx < outputs.size(); idx++) {
    cv::Mat output0 = cv::Mat(cv::Size(100, 6), CV_32F, outputs[idx]);  //[bs,116,8400]=>[bs,8400,116]

    int box_cnt = outputs_size[idx] / step;
    float* pdata = (float*)output0.data;

    for (int box_idx = 0; box_idx < box_cnt; box_idx++) {
      int class_idx = pdata[box_idx * step];
      float score = pdata[box_idx * step + 1];
      float x = pdata[box_idx * step + 2];
      float y = pdata[box_idx * step + 3];
      float w = pdata[box_idx * step + 4];
      float h = pdata[box_idx * step + 5];

      std::cout << class_idx << ", " << score << ", " << x << ", " << y << ", " << w << ", " << h << std::endl;
    }

    // for (int r = 0; r < rows; ++r) {    //stride
    //   cv::Mat scores(1, class_num, CV_32F, pdata + 4);
    //   cv::Point classIdPoint;
    //   double max_class_socre;
    //   cv::minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
    //   max_class_socre = (float)max_class_socre;
    //   if (max_class_socre >= mdl_ox->cfg->draw_threshold) {

    //     //rect [x,y,w,h]
    //     float x = pdata[0] / img_blob.scale;
    //     float y = pdata[1] / img_blob.scale;
    //     int w = pdata[2] / img_blob.scale;
    //     int h = pdata[3] / img_blob.scale;
    //     // std::cout << "score:" << max_class_socre << " idx:" << classIdPoint.x << " x:" << x << " y: " << y << " w: " << w << " h:" << h << std::endl;

    //     int left = std::max(int(x - 0.5 * w + 0.5), 0);
    //     int top = std::max(int(y - 0.5 * h + 0.5), 0);
    //     if (left + w >= img_blob.ori_im_shape[1]) { w = img_blob.ori_im_shape[1] - left - 1;}
    //     if (top + h >= img_blob.ori_im_shape[0]) { h = img_blob.ori_im_shape[0] - top - 1;}
    //     labels.push_back(classIdPoint.x);
    //     confidences.push_back(max_class_socre);
    //     boxs.push_back(cv::Rect(left, top, w, h));
    //   }
    //   pdata += net_width;
    // }
  }

  // std::vector<int> indices;
  // cv::dnn::NMSBoxes(boxs, confidences, mdl_ox->cfg->draw_threshold, mdl_ox->cfg->nms_threshold, indices, 100);
  // for (auto& idx : indices) {
  //   DetInfo* di = new DetInfo(confidences[idx], labels[idx], boxs[idx], mdl_ox->cfg->label_list[labels[idx]]);
  //   det_infos.emplace_back(di);
  // }

  return ret;
}
}